// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Author: Markus Gutschke

#include <arpa/inet.h>
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>

#include "check.h"
#include "connection.h"
#include "hashmap.h"
#include "ijs.h"
#include "stringprintf.h"

#define NOCHECK(x) do { if (x) { } } while (0)


enum IJSCmd {
  IJS_CMD_ACK, IJS_CMD_NAK, IJS_CMD_PING, IJS_CMD_PONG, IJS_CMD_OPEN,
  IJS_CMD_CLOSE, IJS_CMD_BEGIN_JOB, IJS_CMD_END_JOB, IJS_CMD_CANCEL_JOB,
  IJS_CMD_QUERY_STATUS, IJS_CMD_LIST_PARAMS, IJS_CMD_ENUM_PARAM,
  IJS_CMD_SET_PARAM, IJS_CMD_GET_PARAM, IJS_CMD_BEGIN_PAGE,
  IJS_CMD_SEND_DATA_BLOCK, IJS_CMD_END_PAGE, IJS_CMD_EXIT,
};

enum IJSErr {
  IJS_EIO = -2, IJS_EPROTO = -3, IJS_ERANGE = -4, IJS_EINTERNAL = -5,
  IJS_ENYI = -6, IJS_ESYNTAX = -7, IJS_ECOLORSPACE = -8, IJS_EUNKPARAM = -9,
  IJS_EJOBID = -10, IJS_ETOOMANYJOBS = -11, IJS_EBUF = -12
};

struct IJS {
  struct Connection in, out, filter;
  pid_t             pid;
  char              *buf;
  size_t            sz;
  uint32_t          seq;
  void              (*consume)(struct IJS *ijs, const char *buf, size_t sz,
                               void *arg);
  void              *arg;
};


static char *recvData(struct IJS *ijs, enum IJSCmd *ack, uint32_t *sz) {
  uint32_t header[2];
  ssize_t rc = connReadAll(&ijs->in, (char *)header, sizeof(header));
  if (!rc) {
    if (sz) {
      *sz = 0;
    }
    return NULL;
  }
  check(rc == sizeof(header));
  if (ack) {
    *ack = ntohl(header[0]);
  }
  uint32_t len = ntohl(header[1]);
  check(len >= sizeof(header));
  len -= sizeof(header);
  if (sz) {
    *sz = len;
  }
  if (len > 0) {
    if (ijs->sz < len) {
      check(ijs->buf = realloc(ijs->buf, 1+(ijs->sz = len)));
    }
    check(connReadAll(&ijs->in, ijs->buf, len) == len);
    ijs->buf[len] = '\000';
    return ijs->buf;
  } else {
    return "";
  }
}

static void consumeAll(struct IJS *ijs) {
  char buf[8192];
  ssize_t rc;
  while ((rc = connReadTMO(&ijs->filter, buf, sizeof(buf), 0)) > 0) {
    ijs->consume(ijs, buf, rc, ijs->arg);
  }
}

static char *sendCommand(struct IJS *ijs, enum IJSCmd cmd, uint32_t *ret_size,
                         uint32_t *err, const char *buf0, uint32_t sz0,
                         const char *buf1, uint32_t sz1) {
  if (err) {
    *err = 0;
  }
  uint32_t header[] = { htonl(cmd), htonl(sz0 + 2*sizeof(uint32_t)) };
  struct iovec io[] = {
    { .iov_base = header,       .iov_len = sizeof(header) },
    { .iov_base = (char *)buf0, .iov_len = sz0            },
    { .iov_base = (char *)buf1, .iov_len = sz1            }
  };
  for (int n = 0; n < sizeof(io)/sizeof(*io); ) {
    consumeAll(ijs);
    ssize_t sz = connWriteVTMO(&ijs->out, io+n, sizeof(io)/sizeof(*io) - n, 0);
    if (sz < 0 && errno == EAGAIN) {
      struct pollfd pfds[] = { { .fd = ijs->out.fd, .events = POLLOUT },
                               { .fd = ijs->in.fd,  .events = POLLIN  } };
      TEMP_FAILURE_RETRY(poll(pfds, 2, -1));
      continue;
    }
    check(sz > 0);
    while (sz > 0 || (n < sizeof(io)/sizeof(*io) && !io[n].iov_len)) {
      if (sz >= io[n].iov_len) {
        sz -= io[n].iov_len;
        ++n;
      } else {
        io[n].iov_base += sz;
        io[n].iov_len  -= sz;
        break;
      }
    }
  }
  char *data;
  enum IJSCmd ack;
  uint32_t sz;
  switch (cmd) {
  case IJS_CMD_PING:
    check(data = recvData(ijs, &ack, &sz));
    check(ack == IJS_CMD_PONG);
    check(sz == sizeof(uint32_t));
    uint32_t version = ntohl(*(uint32_t *)data);
    check(version >= 30);
    break;
  case IJS_CMD_QUERY_STATUS:
  case IJS_CMD_LIST_PARAMS:
  case IJS_CMD_ENUM_PARAM:
  case IJS_CMD_GET_PARAM:
    check(data = recvData(ijs, &ack, &sz));
    check(ack == IJS_CMD_ACK || ack == IJS_CMD_NAK);
    if (ack == IJS_CMD_NAK) {
      check(sz == sizeof(uint32_t));
      if (err) {
        *err = ntohl(*(uint32_t *)data);
      }
    }
    break;
  default:
    check(data = recvData(ijs, &ack, &sz));
    check(ack == IJS_CMD_ACK || ack == IJS_CMD_NAK);
    if (ack == IJS_CMD_ACK) {
      check(sz == 0);
    } else {
      check(sz == sizeof(uint32_t));
      if (err) {
        *err = ntohl(*(uint32_t *)data);
      }
    }
    break;
  }
  if (ret_size) {
    *ret_size = sz;
  }
  return data;
}

static char *sendInteger(struct IJS *ijs, enum IJSCmd cmd, uint32_t *ret_size,
                         uint32_t *err, uint32_t arg) {
  arg = htonl(arg);
  return sendCommand(ijs, cmd, ret_size, err, (char *)&arg, sizeof(arg),
                     NULL, 0);
}

struct IJS *newIJS(const char *server,
                   void (*consume)(struct IJS *ijs, const char *buf, size_t sz,
                                   void *arg),
                   void *arg) {
  struct IJS *ijs;
  check(ijs = malloc(sizeof(struct IJS)));
  initIJS(ijs, server, consume, arg);
  return ijs;
}

void initIJS(struct IJS *ijs, const char *server,
             void (*consume)(struct IJS *ijs, const char *buf, size_t sz,
                             void *arg), void *arg) {
  int fds[8];
  check(!pipe(fds) && !pipe(fds + 2) && !pipe(fds + 4) && !pipe(fds + 6));
  for (int i = 0; i < sizeof(fds)/sizeof(*fds); ++i) {
    check(fds[i] > 2);
  }
  ijs->pid = fork();
  check(ijs->pid >= 0);
  if (ijs->pid == 0) {
    if (dup2(fds[5], 1) != 1 || dup2(fds[5], 2) != 2) {
    child_failed:;
      const char *err = strerror(errno);
      TEMP_FAILURE_RETRY(write(fds[5], err, strlen(err)));
      _exit(1);
    }
    check(!fcntl(fds[5], F_SETFD, FD_CLOEXEC | fcntl(fds[5], F_GETFD, 0)));
    DIR *dir = opendir("/proc/self/fd");
    if (dir == 0) {
      for (int fd = sysconf(_SC_OPEN_MAX); --fd > 2; ) {
        if (fd != fds[1] && fd != fds[2] && fd != fds[5] && fd != fds[7]) {
          TEMP_FAILURE_RETRY(close(fd));
        }
      }
    } else {
      struct dirent de, *res;
      while (!readdir_r(dir, &de, &res) && res) {
        if (res->d_name[0] < '0')
          continue;
        int fd = atoi(res->d_name);
        if (fd > 2 && fd != fds[1] && fd != fds[2] && fd != fds[5] &&
            fd != fds[7] && fd != dirfd(dir)) {
          TEMP_FAILURE_RETRY(close(fd));
        }
      }
      check(!closedir(dir));
    }
    int devnull;
    check ((devnull = open("/dev/null", O_RDWR)) > 2);
    check(dup2(fds[2],  0) == 0);
    check(dup2(fds[1],  1) == 1);
    check(dup2(devnull, 2) == 2);
    check(dup2(fds[7],  3) == 3);
    if (fds[2]  != 3) check(!TEMP_FAILURE_RETRY(close(fds[2])));
    if (fds[1]  != 3) check(!TEMP_FAILURE_RETRY(close(fds[1])));
    if (devnull != 3) check(!TEMP_FAILURE_RETRY(close(devnull)));
    if (fds[7]  != 3) check(!TEMP_FAILURE_RETRY(close(fds[7])));
    execlp(server, server, NULL);
    goto child_failed;
  }
  check(!TEMP_FAILURE_RETRY(close(fds[1])));
  check(!TEMP_FAILURE_RETRY(close(fds[2])));
  check(!TEMP_FAILURE_RETRY(close(fds[5])));
  check(!TEMP_FAILURE_RETRY(close(fds[7])));
  for (int failed = 0;;) {
    char buf[40];
    ssize_t sz = TEMP_FAILURE_RETRY(read(fds[4], buf, sizeof(buf)));
    check(sz >= 0);
    if (!sz) {
      if (failed) {
        TEMP_FAILURE_RETRY(write(2, "\n", 1));
        _exit(1);
      }
      break;
    } else {
      failed = 1;
      TEMP_FAILURE_RETRY(write(2, buf, sz));
    }
  }
  check(!TEMP_FAILURE_RETRY(close(fds[4])));
  check(!fcntl(fds[3], F_SETFL, (long)O_NONBLOCK));
  check(!fcntl(fds[6], F_SETFL, (long)O_NONBLOCK));
  initConnection(&ijs->in,     fds[0]);
  initConnection(&ijs->out,    fds[3]);
  initConnection(&ijs->filter, fds[6]);
  ijs->buf     = NULL;
  ijs->sz      = 0;
  ijs->seq     = 0;
  ijs->consume = consume;
  ijs->arg     = arg;
  check(connWrite(&ijs->out, "IJS\n\252v1\n", 8) == 8);
  char buf[8];
  check(connReadAll(&ijs->in, buf, 8) == 8);
  check(!memcmp(buf, "IJS\n\253v1\n", 8));
  uint32_t err;
  sendInteger(ijs, IJS_CMD_PING, NULL, &err, 30);
  check(!err);
}

void destroyIJS(struct IJS *ijs) {
  if (ijs) {
    sendCommand(ijs, IJS_CMD_EXIT, NULL, NULL, NULL, 0, NULL, 0);
    int processStatus;
    check(waitpid(ijs->pid, &processStatus, 0) == ijs->pid);
    check(WIFEXITED(processStatus));
    consumeAll(ijs);
    free(ijs->buf);
    destroyConnection(&ijs->in);
    destroyConnection(&ijs->out);
    destroyConnection(&ijs->filter);
  }
}

void deleteIJS(struct IJS *ijs) {
  destroyIJS(ijs);
  free(ijs);
}

void ijsOpen(struct IJS *ijs) {
  uint32_t err;
  sendCommand(ijs, IJS_CMD_OPEN, NULL, &err, NULL, 0, NULL, 0);
  check(!err);
}

void ijsClose(struct IJS *ijs) {
  uint32_t err;
  sendCommand(ijs, IJS_CMD_CLOSE, NULL, &err, NULL, 0, NULL, 0);
  check(!err);
}

uint32_t ijsBeginJob(struct IJS *ijs) {
  uint32_t err;
  sendInteger(ijs, IJS_CMD_BEGIN_JOB, NULL, &err, ++ijs->seq);
  check(!err);
  check(!ijsSetParam(ijs, ijs->seq, "OutputFD", "3", -1));
  check(!ijsQueryStatus(ijs, ijs->seq));
  return ijs->seq;
}

void ijsEndJob(struct IJS *ijs, uint32_t job) {
  uint32_t err;
  sendInteger(ijs, IJS_CMD_END_JOB, NULL, &err, job);
  check(!err);
}

uint32_t ijsQueryStatus(struct IJS *ijs, uint32_t job) {
  uint32_t ret_size, err, status;
  char *data = sendInteger(ijs, IJS_CMD_QUERY_STATUS, &ret_size, &err, job);
  check(!err);
  check(ret_size == sizeof(uint32_t));
  memcpy(&status, data, ret_size);
  return status;
}

int ijsSetParam(struct IJS *ijs, uint32_t job, const char *param,
                const char *value, ssize_t sz) {
  // The wire format is actually very different from the one documented in
  // the IJS specification. It turns out, the specification does not match
  // existing implementations at all.
  size_t paramLen = strlen(param);
  if (sz < 0) {
    sz = strlen(value);
  }
  char *args;
  check(args = malloc(2*sizeof(uint32_t) + paramLen + sz + 1));
  ((uint32_t *)args)[0] = htonl(job);
  ((uint32_t *)args)[1] = htonl(paramLen + 1 + sz);
  memcpy(strcpy(args + 2*sizeof(uint32_t), param) + paramLen + 1, value, sz);
  uint32_t ret_size, err;
  sendCommand(ijs, IJS_CMD_SET_PARAM, &ret_size, &err, args,
              2*sizeof(uint32_t) + paramLen + 1 + sz, NULL, 0);
  free(args);
  return err;
}

int ijsSetParamF(struct IJS *ijs, uint32_t job, const char *param,
                 const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  char *s = vStringPrintf(NULL, fmt, ap);
  va_end(ap);
  int rc = ijsSetParam(ijs, job, param, s, -1);
  free(s);
  return rc;
}

const char *ijsGetParam(struct IJS *ijs, uint32_t *ret_size, uint32_t job,
                        const char *param) {
  size_t paramLen = strlen(param);
  char *args;
  check(args = malloc(sizeof(uint32_t) + paramLen + 1));
  *(uint32_t *)args = htonl(job);
  strcpy(args + sizeof(uint32_t), param);
  uint32_t err;
  char *data = sendCommand(ijs, IJS_CMD_GET_PARAM, ret_size, &err, args,
                           sizeof(uint32_t) + paramLen + 1, NULL, 0);
  free(args);
  return err ? NULL : data;
}

static void destroyEnumParamEntry(void *arg, char *key, char *value) {
  free(key);
}

struct HashMap *ijsEnumParam(struct IJS *ijs, uint32_t job, const char *param){
  size_t len = strlen(param);
  char *args, *data;
  check(args = malloc(sizeof(uint32_t) + len + 1));
  *(uint32_t *)args = htonl(job);
  strcpy(args + sizeof(uint32_t), param);
  uint32_t ret_size, err;
  data = sendCommand(ijs, IJS_CMD_ENUM_PARAM, &ret_size, &err, args,
                     sizeof(uint32_t) + len + 1, NULL, 0);
  free(args);
  if (err) {
    return NULL;
  } else {
    struct HashMap *map = newHashMap(destroyEnumParamEntry, NULL);
    for (char *ptr = data;;) {
      while (*ptr == ',') {
        ++ptr;
      }
      if (!*ptr) {
        break;
      }
      len = strcspn(ptr, ",");
      char *key;
      check(key = malloc(len + 1));
      memcpy(key, ptr, len);
      key[len] = '\000';
      addToHashMap(map, key, key);
      ptr += len;
    }
    return map;
  }
}

static void destroyListParamsEntry(void *arg, char *key, char *value) {
  if (key == arg) {
    free(key);
  }
  deleteHashMap((struct HashMap *)value);
}

struct HashMap *ijsListParams(struct IJS *ijs, uint32_t job) {
  uint32_t ret_size, err;
  char *data = sendInteger(ijs, IJS_CMD_LIST_PARAMS, &ret_size, &err, job);
  check(!err);
  while (*data == ',') {
    ++data;
  }
  if (!*data) {
    return newHashMap(NULL, NULL);
  }
  check(data = strdup(data));
  struct HashMap *map = newHashMap(destroyListParamsEntry, data);
  for (char *ptr = data; *ptr;) {
    size_t len = strcspn(ptr, ",");
    int done = !ptr[len];
    ptr[len] = '\000';
    addToHashMap(map, ptr, (char *)ijsEnumParam(ijs, job, ptr));
    if (done) {
      break;
    }
    ptr += len;
    while (*++ptr == ',') { }
  }
  return map;
}

void ijsBeginPage(struct IJS *ijs, uint32_t job) {
  uint32_t err;
  sendCommand(ijs, IJS_CMD_BEGIN_PAGE, NULL, &err, NULL, 0, NULL, 0);
  check(!err);
}

void ijsEndPage(struct IJS *ijs, uint32_t job) {
  uint32_t err;
  sendCommand(ijs, IJS_CMD_END_PAGE, NULL, &err, NULL, 0, NULL, 0);
  check(!err);
}

void ijsSendData(struct IJS *ijs, uint32_t job, const char *data,
                 uint32_t sz) {
  uint32_t err, args[] = { htonl(job), htonl(sz) };
  sendCommand(ijs, IJS_CMD_SEND_DATA_BLOCK, NULL, &err, (const char *)args,
              sizeof(args), data, sz);
  check(!err);
}
